import requests
import time
from datetime import datetime, time as dtime, timedelta
import json
import os
import math
import sys
from typing import List, Tuple, Dict
from collections import defaultdict
import random

class TimeScheduler:
    """时间调度管理器"""
    def __init__(self):
        self.interval_minutes = 10
    
    def get_next_interval(self) -> datetime:
        """计算下一个精确的10分钟间隔时间点"""
        current_time = datetime.now()
        # 将当前时间规整到最近的10分钟间隔
        minutes = (current_time.minute // self.interval_minutes) * self.interval_minutes
        next_time = current_time.replace(minute=minutes, second=0, microsecond=0)
        
        # 如果当前时间已经过了这个间隔，就取下一个间隔
        if current_time > next_time:
            next_time = next_time + timedelta(minutes=self.interval_minutes)
            
        return next_time

class GeoGrid:
    """地理网格划分"""
    EARTH_RADIUS = 6371  # 地球半径，单位公里
    
    def __init__(self, min_lng: float, min_lat: float, max_lng: float, max_lat: float, max_diagonal: float = 2.0):
        self.min_lng = min_lng
        self.min_lat = min_lat
        self.max_lng = max_lng
        self.max_lat = max_lat
        self.max_diagonal = max_diagonal
        self.grid_cells = self.create_grid()
    
    def calculate_distance(self, lng1: float, lat1: float, lng2: float, lat2: float) -> float:
        """计算两点间的距离（公里）"""
        rad_lat1 = math.radians(lat1)
        rad_lat2 = math.radians(lat2)
        delta_lat = math.radians(lat2 - lat1)
        delta_lng = math.radians(lng2 - lng1)
        
        a = (math.sin(delta_lat/2) * math.sin(delta_lat/2) +
             math.cos(rad_lat1) * math.cos(rad_lat2) *
             math.sin(delta_lng/2) * math.sin(delta_lng/2))
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
        return self.EARTH_RADIUS * c
    
    def create_grid(self) -> List[Tuple[float, float, float, float]]:
        """将大矩形区域划分为小网格"""
        diagonal = self.calculate_distance(self.min_lng, self.min_lat, self.max_lng, self.max_lat)
        n = math.ceil(diagonal / self.max_diagonal)
        
        lng_step = (self.max_lng - self.min_lng) / n
        lat_step = (self.max_lat - self.min_lat) / n
        
        grid_cells = []
        for i in range(n):
            for j in range(n):
                cell_min_lng = self.min_lng + i * lng_step
                cell_min_lat = self.min_lat + j * lat_step
                cell_max_lng = min(cell_min_lng + lng_step, self.max_lng)
                cell_max_lat = min(cell_min_lat + lat_step, self.max_lat)
                grid_cells.append((cell_min_lng, cell_min_lat, cell_max_lng, cell_max_lat))
        
        return grid_cells

class AKManager:
    """AK管理器：处理多个百度地图API密钥的使用和配额控制"""
    def __init__(self, ak_list: List[str]):
        self.ak_list = ak_list
        self.daily_quota = 18  # 每个AK每天的配额次数
        self.usage_count = defaultdict(int)  # 记录每个AK的使用次数
        self.error_count = defaultdict(int)  # 记录每个AK的错误次数
        self.last_reset = datetime.now().date()
        self.current_ak_index = 0
        self.ak_status = {ak: True for ak in ak_list}  # True表示AK可用
        
    def reset_daily_count(self):
        """每天重置使用次数"""
        current_date = datetime.now().date()
        if current_date > self.last_reset:
            self.usage_count.clear()
            self.error_count.clear()
            self.ak_status = {ak: True for ak in self.ak_list}
            self.last_reset = current_date
            print("已重置所有AK的使用计数")

    def get_next_available_ak(self) -> str:
        """获取下一个可用的AK"""
        self.reset_daily_count()
        
        available_aks = [ak for ak in self.ak_list if self.ak_status[ak]]
        if not available_aks:
            raise Exception("所有AK都已达到配额限制或出现错误")
        
        valid_aks = [ak for ak in available_aks if self.usage_count[ak] < self.daily_quota]
        if valid_aks:
            selected_ak = random.choice(valid_aks)
            self.usage_count[selected_ak] += 1
            return selected_ak
            
        selected_ak = min(available_aks, key=lambda x: self.error_count[x])
        self.usage_count[selected_ak] += 1
        return selected_ak

    def report_error(self, ak: str, error_type: str):
        """报告AK使用时的错误"""
        self.error_count[ak] += 1
        print(f"AK {ak[:8]}... 发生错误: {error_type}")
        
        if error_type == "配额超限" or self.error_count[ak] >= 3:
            self.ak_status[ak] = False
            print(f"AK {ak[:8]}... 已被标记为不可用")

    def get_status_report(self) -> Dict:
        """获取所有AK的状态报告"""
        return {
            'total_aks': len(self.ak_list),
            'available_aks': sum(1 for ak in self.ak_status if self.ak_status[ak]),
            'usage_counts': dict(self.usage_count),
            'error_counts': dict(self.error_count),
            'ak_status': self.ak_status
        }

def ensure_directory(directory_path: str) -> bool:
    """确保目录存在，如果不存在则尝试创建"""
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        full_path = os.path.join(script_dir, directory_path)
        
        if not os.path.exists(full_path):
            os.makedirs(full_path)
            print(f"成功创建数据存储目录: {full_path}")
        return True
    except Exception as e:
        print(f"创建目录失败: {str(e)}")
        print(f"尝试使用的路径: {full_path}")
        return False

class TrafficMonitor:
    def __init__(self, ak_list: List[str]):
        self.url = "https://api.map.baidu.com/traffic/v1/bound"
        self.ak_manager = AKManager(ak_list)
        self.grid = GeoGrid(106.4847586, 29.5319042, 106.5991668, 29.5792822)
        self.scheduler = TimeScheduler()
        
        if not ensure_directory('traffic_data'):
            print("错误：无法创建数据存储目录，程序将退出")
            sys.exit(1)
    
    def query_traffic_for_cell(self, cell: Tuple[float, float, float, float]) -> dict:
        """查询单个网格的交通状况"""
        max_retries = 3
        for attempt in range(max_retries):
            try:
                ak = self.ak_manager.get_next_available_ak()
                params = {
                    "bounds": f"{cell[1]},{cell[0]};{cell[3]},{cell[2]}",
                    "road_grade": "0",
                    "coord_type_input": "gcj02",
                    "coord_type_output": "gcj02",
                    "ak": ak
                }
                
                response = requests.get(url=self.url, params=params)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('status') == 302:
                        self.ak_manager.report_error(ak, "配额超限")
                        continue
                        
                    data['query_bounds'] = params['bounds']
                    data['query_timestamp'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    return data
                else:
                    self.ak_manager.report_error(ak, f"HTTP错误 {response.status_code}")
                    
            except Exception as e:
                self.ak_manager.report_error(ak, str(e))
                
            if attempt == max_retries - 1:
                print(f"网格查询失败，已重试{max_retries}次")
                return None
                
            time.sleep(1)
    
    def query_all_traffic(self) -> List[dict]:
        """查询所有网格的交通状况"""
        all_results = []
        for cell in self.grid.grid_cells:
            result = self.query_traffic_for_cell(cell)
            if result:
                all_results.append(result)
            time.sleep(0.5)
        
        status = self.ak_manager.get_status_report()
        print("\nAK状态报告:")
        print(f"可用AK数量: {status['available_aks']}/{status['total_aks']}")
        print("使用次数:", {k[:8]: v for k, v in status['usage_counts'].items()})
        
        return all_results
    
    def save_results(self, results: List[dict]):
        """保存查询结果"""
        if not results:
            return
            
        try:
            current_time = datetime.now()
            script_dir = os.path.dirname(os.path.abspath(__file__))
            filename = os.path.join(
                script_dir,
                'traffic_data',
                f"traffic_{current_time.strftime('%Y%m%d_%H%M')}.json"
            )
            
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            
            print(f"数据已保存至: {filename}")
            
        except Exception as e:
            print(f"保存数据时发生错误: {str(e)}")

def main():
    # 系统信息输出
    print("系统信息：")
    print(f"当前工作目录: {os.getcwd()}")
    print(f"脚本位置: {os.path.abspath(__file__)}")
    
    # API密钥列表
    ak_list = [
        "*******",
        "*******"
    ]
    
    monitor = TrafficMonitor(ak_list)
    
    print("\n交通监控程序已启动...")
    print(f"总网格数量: {len(monitor.grid.grid_cells)}")
    print(f"可用AK数量: {len(ak_list)}")
    print("监控模式：全天候监控，每10分钟一次")
    
    try:
        # 等待到下一个10分钟间隔开始
        next_interval = monitor.scheduler.get_next_interval()
        wait_seconds = (next_interval - datetime.now()).total_seconds()
        if wait_seconds > 0:
            print(f"\n等待至下一个整10分钟: {next_interval.strftime('%Y-%m-%d %H:%M:%S')}")
            print(f"等待时间: {wait_seconds:.2f}秒")
            time.sleep(wait_seconds)
        
        while True:
            start_time = time.time()
            current_time = datetime.now()
            
            print(f"\n开始新一轮查询: {current_time.strftime('%Y-%m-%d %H:%M:%S')}")
            results = monitor.query_all_traffic()
            monitor.save_results(results)
            
            # 计算实际处理时间
            processing_time = time.time() - start_time
            
            # 计算到下一个10分钟间隔的等待时间
            next_query_time = monitor.scheduler.get_next_interval()
            wait_seconds = (next_query_time - datetime.now()).total_seconds()
            wait_seconds = max(0, wait_seconds)
            
            print(f"处理耗时: {processing_time:.2f}秒")
            print(f"等待{wait_seconds:.2f}秒后进行下一轮查询...")
            print(f"下次查询时间: {next_query_time.strftime('%Y-%m-%d %H:%M:%S')}")
            
            time.sleep(wait_seconds)
                
    except KeyboardInterrupt:
        print("\n程序已停止")

if __name__ == "__main__":
    main()
